import backtrader as bt
import numpy as np
import pandas as pd
from typing import Dict, Any, Optional

class MeanReversionStrategy(bt.Strategy):
    """
    均值回归策略
    
    策略逻辑：
    1. 使用布林带识别超买超卖
    2. RSI指标确认反转信号
    3. 在价格触及下轨且RSI < 30时买入
    4. 在价格触及上轨且RSI > 70时卖出
    5. 设置止损和止盈
    """
    
    params = (
        ('bb_period', 20),      # 布林带周期
        ('bb_dev', 2),          # 布林带标准差倍数
        ('rsi_period', 14),     # RSI周期
        ('rsi_oversold', 30),   # RSI超卖阈值
        ('rsi_overbought', 70), # RSI超买阈值
        ('stop_loss', 0.05),    # 止损比例
        ('take_profit', 0.15),  # 止盈比例
        ('position_size', 0.1), # 仓位比例
    )
    
    def __init__(self):
        # 技术指标
        self.bb = bt.indicators.BollingerBands(
            self.data.close, 
            period=self.params.bb_period,
            devfactor=self.params.bb_dev
        )
        self.rsi = bt.indicators.RSI(
            self.data.close, 
            period=self.params.rsi_period
        )
        
        # 交易状态
        self.order = None
        self.stop_order = None
        self.take_profit_order = None
        
    def next(self):
        if self.order:
            return
            
        if not self.position:
            # 寻找买入机会
            if (self.data.close[0] <= self.bb.lines.bot[0] and 
                self.rsi[0] < self.params.rsi_oversold):
                
                # 计算买入数量
                cash = self.broker.getcash()
                price = self.data.close[0]
                size = int((cash * self.params.position_size) / price)
                
                if size > 0:
                    self.order = self.buy(size=size)
                    print(f'买入信号: 价格={price:.2f}, RSI={self.rsi[0]:.2f}')
        
        else:
            # 寻找卖出机会
            if (self.data.close[0] >= self.bb.lines.top[0] and 
                self.rsi[0] > self.params.rsi_overbought):
                
                self.order = self.sell(size=self.position.size)
                print(f'卖出信号: 价格={self.data.close[0]:.2f}, RSI={self.rsi[0]:.2f}')
    
    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            return
            
        if order.status in [order.Completed]:
            if order.isbuy():
                # 设置止损和止盈
                price = order.executed.price
                stop_price = price * (1 - self.params.stop_loss)
                take_profit_price = price * (1 + self.params.take_profit)
                
                self.stop_order = self.sell(
                    size=order.executed.size,
                    exectype=bt.Order.Stop,
                    price=stop_price
                )
                
                self.take_profit_order = self.sell(
                    size=order.executed.size,
                    exectype=bt.Order.Limit,
                    price=take_profit_price
                )
                
                print(f'买入完成: 价格={price:.2f}, 数量={order.executed.size}')
            else:
                # 取消其他订单
                if self.stop_order:
                    self.cancel(self.stop_order)
                if self.take_profit_order:
                    self.cancel(self.take_profit_order)
                    
                print(f'卖出完成: 价格={order.executed.price:.2f}, 数量={order.executed.size}')
        
        self.order = None

def get_strategy_params() -> Dict[str, Any]:
    """获取策略参数"""
    return {
        'name': '均值回归策略',
        'description': '基于布林带和RSI的均值回归策略，在超卖时买入，超买时卖出',
        'parameters': {
            'bb_period': {'type': 'int', 'default': 20, 'min': 10, 'max': 50, 'description': '布林带周期'},
            'bb_dev': {'type': 'float', 'default': 2.0, 'min': 1.5, 'max': 3.0, 'description': '布林带标准差倍数'},
            'rsi_period': {'type': 'int', 'default': 14, 'min': 10, 'max': 21, 'description': 'RSI周期'},
            'rsi_oversold': {'type': 'int', 'default': 30, 'min': 20, 'max': 40, 'description': 'RSI超卖阈值'},
            'rsi_overbought': {'type': 'int', 'default': 70, 'min': 60, 'max': 80, 'description': 'RSI超买阈值'},
            'stop_loss': {'type': 'float', 'default': 0.05, 'min': 0.02, 'max': 0.10, 'description': '止损比例'},
            'take_profit': {'type': 'float', 'default': 0.15, 'min': 0.08, 'max': 0.25, 'description': '止盈比例'},
            'position_size': {'type': 'float', 'default': 0.1, 'min': 0.05, 'max': 0.3, 'description': '仓位比例'}
        }
    } 